/*
 * Copyright (C) 2007 Eskil Bylund
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

using System;
using System.Collections.Generic;
using System.IO;
using System.Security.Cryptography;

namespace DCSharp.Hashing
{
	/// <summary>
	/// Represents a Merkle Hash Tree.
	/// </summary>
	/// <remarks>
	/// <para>The leaves will be created from the number of bytes specified by
	/// SegmentSize. If a block size is specified the leaves will be combined
	/// until each leaf represents that many bytes of the original data. Using
	/// a large block size is useful on large files to keep the leaf count down.
	/// </para>
	/// <para>The specification for Merkle Hash Trees are available at
	/// http://open-content.net/specs/draft-jchapweske-thex-02.html</para>
	/// <para>http://en.wikipedia.org/wiki/Tree_data_structure</para>
	/// </remarks> 
	public class HashTree
	{
		/// <summary>
		/// The size of each file segment used to create the tree.
		/// </summary>
		public const int SegmentSize = 1024;

		private const byte InternalHashPrefix = 0x01;
		private const byte LeafHashPrefix = 0x00;

		private HashAlgorithm algorithm;
		private List<byte[]> leaves;
		private List<byte[]> segments;

		private long blockSize;
		private long fileSize;
		private byte[] root;
		private byte[] buffer;

		/// <summary>
		/// Constructs a new HashTree with the given hash algorithm.
		/// </summary>
		/// <param name="algorithm">The hash algorithm.</param>
		public HashTree(HashAlgorithm algorithm) : this(algorithm, SegmentSize)
		{
		}

		/// <summary>
		/// Constructs a new HashTree with the given hash algorithm and block size.
		/// </summary>
		/// <param name="algorithm">The hash algorithm.</param>
		/// <param name="blockSize">The block size.</param>
		public HashTree(HashAlgorithm algorithm, long blockSize) :
			this(algorithm, blockSize, 0, null, null)
		{
		}

		/// <summary>
		/// Constructs a HashTree with the given hash algorithm, block size and
		/// file size, with leaves read from the stream.
		/// </summary>
		/// <param name="algorithm">The hash algorithm.</param>
		/// <param name="blockSize">The block size.</param>
		/// <param name="fileSize">The file size.</param>
		/// <param name="stream">The stream containing the leaves.</param>
		/// <param name="root">Optional root hash.</param>
		public HashTree(HashAlgorithm algorithm, long blockSize, long fileSize,
			Stream stream, byte[] root)
		{
			if (algorithm == null)
			{
				throw new ArgumentNullException("algorithm");
			}
			if ((blockSize / SegmentSize) % 2 != 0 && blockSize != SegmentSize)
			{
				throw new ArgumentException(
					"The block size must be a multiple of the segment size.",
					"blockSize");
			}

			this.algorithm = algorithm;
			this.blockSize = blockSize;
			this.fileSize = fileSize;
			this.root = root;

			leaves = new List<byte[]>();
			segments = new List<byte[]>();
			buffer = new byte[SegmentSize + 1];

			if (stream != null)
			{
				LoadLeaves(stream);
			}
		}

		#region Properties

		/// <summary>
		/// Gets the size of the hash computed by the hash algorithm.
		/// </summary>
		/// <value>The size of the hash computed by the hash algorithm.</value>
		public int HashSize
		{
			get { return algorithm.HashSize; }
		}

		/// <summary>
		/// Gets the block size.
		/// </summary>
		/// <value>The block size.</value>
		public long BlockSize
		{
			get { return blockSize; }
		}

		/// <summary>
		/// Gets the size of the file that the tree was generated from.
		/// </summary>
		/// <value>The size of the file that the tree was generated from.</value>
		public long FileSize
		{
			get { return fileSize; }
		}

		/// <summary>
		/// Gets the number of leaves in the tree.
		/// </summary>
		/// <value>The number of leaves in the tree.</value>
		public int LeafCount
		{
			get { return leaves.Count; }
		}

		/// <summary>
		/// Gets the leaves of the tree.
		/// </summary>
		/// <value>The leaves of the tree.</value>
		public byte[][] Leaves
		{
			get { return leaves.ToArray(); }
		}

		/// <summary>
		/// Gets the previously computed root hash.
		/// </summary>
		/// <value>The previously computed root hash.</value>
		public byte[] Root
		{
			get { return root; }
		}

		public byte[] this[int index]
		{
			get { return leaves[index]; }
		}

		#endregion

		#region Methods

		/// <summary>
		/// Updates the tree with data.
		/// </summary>
		/// <param name="input">The byte[] with the data.</param>
		public void Update(byte[] input)
		{
			Update(input, 0, input.Length);
		}

		/// <summary>
		/// Updates the tree with data.
		/// </summary>
		/// <param name="input">The byte[] with the data.</param>
		/// <param name="offset">The position in input at which to begin reading.</param>
		/// <param name="count">
		/// The number of bytes to read. Must be a multiple of SegmentSize,
		/// unless it's the last segment.
		/// </param>
		public void Update(byte[] input, int offset, int count)
		{
			int length = offset + count;
			for (int pos = offset; pos < length; pos += SegmentSize)
			{
				int n = Math.Min(SegmentSize, length - pos);

				buffer[0] = LeafHashPrefix;

				algorithm.Initialize();
				algorithm.TransformBlock(buffer, 0, 1, buffer, 1);
				algorithm.TransformFinalBlock(input, pos, n);

				if (blockSize == SegmentSize)
				{
					leaves.Add(algorithm.Hash);
				}
				else
				{
					segments.Add(algorithm.Hash);

					long segmentsPerBlock = blockSize / SegmentSize;
					if (segments.Count == segmentsPerBlock)
					{
						leaves.Add(GetRootHash(algorithm, segments));
						segments.Clear();
					}
				}
				fileSize += n;
			}
		}

		/// <summary>
		/// Updates the tree with data from a stream.
		/// </summary>
		/// <param name="stream">
		/// The stream with the data. The stream length must be a multiple of
		/// SegmentSize, unless the stream contains the last segment.
		/// </param>
		public void Update(Stream stream)
		{
			Update(stream, stream.Length - stream.Position);
		}

		/// <summary>
		/// Updates the tree with data from a stream.
		/// </summary>
		/// <param name="stream">The stream with the data.</param>
		/// <param name="count">
		/// The number of bytes to read. Must be a multiple of SegmentSize,
		/// unless the stream contains the last segment.
		/// </param>
		public void Update(Stream stream, long count)
		{
			byte[] data = new byte[SegmentSize * 8];
			int read;
			do
			{
				read = stream.Read(data, 0, (int)Math.Min(count, data.Length));
				if (read != 0)
				{
					Update(data, 0, read);
					count -= read;
				}
			}
			while (read != 0);
		}

		/// <summary>
		/// Calculates the root hash from the leaves.
		/// </summary>
		/// <returns>A byte[] containing the root hash.</returns>
		public byte[] CalculateRoot()
		{
			if (segments.Count > 0)
			{
				leaves.Add(GetRootHash(algorithm, segments));
				segments.Clear();
			}

			if (leaves.Count == 0)
			{
				root = new byte[1];
				root = algorithm.ComputeHash(root, 0, 1);
			}
			else
			{
				root = GetRootHash(algorithm, leaves);
			} 
			return root;
		}

		/// <summary>
		/// Calculates the block count for a specific file and block size.
		/// </summary>
		/// <param name="fileSize">The file size.</param>
		/// <param name="blockSize">The block size.</param>
		/// <returns>The calculated block count.</returns>
		public static int CalculateBlockCount(long fileSize, long blockSize)
		{
			return Math.Max((int)((fileSize + blockSize - 1) / blockSize), 1);
		}

		/// <summary>
		/// Calculates the block size needed to keep the tree under a specific depth. 
		/// </summary>
		/// <param name="fileSize">The file size.</param>
		/// <param name="maxLevels">The maximum depth of the tree.</param>
		/// <returns>The calculated block size.</returns>
		public static long CalculateBlockSize(long fileSize, int maxLevels)
		{
			long tmp = SegmentSize;
			long maxHashes = 1 << (maxLevels - 1);
			while ((maxHashes * tmp) < fileSize)
			{
				tmp *= 2;
			}
			return tmp;
		}

		/// <summary>
		/// A helper method for checking if two leaves are equal.
		/// </summary>
		/// <param name="leafA">The first leaf.</param>
		/// <param name="leafB">The second leaf.</param>
		/// <returns>True if the leaves are equal; otherwise, false.</returns>
		public static bool LeavesAreEqual(byte[] leafA, byte[] leafB)
		{
			if (leafA.Length != leafB.Length)
			{
				return false;
			}
			for (int j = 0; j < leafA.Length; j++)
			{
				if (leafA[j] != leafB[j])
				{
					return false;
				}
			}
			return true;
		}

		private void LoadLeaves(Stream stream)
		{
			int blockCount = CalculateBlockCount(fileSize, blockSize);
			byte[] buffer = new byte[algorithm.HashSize];

			for (int i = 0; i < blockCount; i++)
			{
				int count = stream.Read(buffer, 0, buffer.Length);
				if (count != 0)
				{
					byte[] leaf = new byte[count];
					Buffer.BlockCopy(buffer, 0, leaf, 0, count);

					leaves.Add(leaf);
				}
				else
				{
					throw new IOException("Unexpected end of file.");
				}
			}
		}

		private byte[] GetRootHash(HashAlgorithm hash, List<byte[]> leaves)
		{
			if (leaves.Count == 0)
			{
				throw new ArgumentException("Nothing to hash");
			}

			Queue<byte[]> leavesToHash = new Queue<byte[]>();
			for (int i = 0; i < leaves.Count; i++)
			{
				leavesToHash.Enqueue(leaves[i]);
			}

			buffer[0] = InternalHashPrefix;

			while (leavesToHash.Count > 1)
			{
				int count = leavesToHash.Count;
				for (int i = 0; i + 1 < count; i += 2)
				{
					byte[] left = leavesToHash.Dequeue();
					byte[] right = leavesToHash.Dequeue();

					algorithm.Initialize();
					algorithm.TransformBlock(buffer, 0, 1, buffer, 1);
					algorithm.TransformBlock(left, 0, left.Length, buffer, 1);
					algorithm.TransformFinalBlock(right, 0, right.Length);

					leavesToHash.Enqueue(algorithm.Hash);
				}
				if (count % 2 != 0)
				{
					leavesToHash.Enqueue(leavesToHash.Dequeue());
				}
			}
			return leavesToHash.Dequeue();
		}

		#endregion
	}
}
